{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Language modelling\n",
    "*Thomas Dooms*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The previous chapter covered how to study bilinear MLPs in a real-world scenario. We also discussed how to decompose interaction matrices into eigenvectors, which are interpretable and causally relevant. We now turn to language models and explain how this method can be leveraged to better understand latent interactions. This is done by leveraging latent feature bases, such as features from an SAE. Our aim is to demonstrate that weight-based interpretability is not some theoretical method that only works in small models but that it can be combined with many other techniques in complex real-world environments. In essence, weight-based interpretability allows us to 'trace' between existing features to better understand how they are formed or what their effect is.\n",
    "\n",
    "The main analysis technique we will discuss here is training two SAEs around the bilinear MLP. We then use the features from the output SAE as output features (just like the digits we chose previously). That way, we can decompose the MLP to which SAE input features interact strongly toward that output. In a sense, this technique finds shallow circuits (but grounded in the weights) in an MLP. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from language import Transformer\n",
    "from sae import Tracer, Visualizer\n",
    "from sae.functions import compute_truncated_eigenvalues\n",
    "\n",
    "import plotly.express as px\n",
    "import torch\n",
    "\n",
    "device = \"cuda:0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up.\n",
    "We generously received quite a bit of compute from Eleuther to train capable bilinear transformers. Here, we will study a GPT2-medium level model (which is probably somewhat undertrained) but can still produce coherent sentences and has non-trivial knowledge about some topics. The model has 16 layers and about 300M parameters. We have trained SAEs around each MLP layer for later analysis. \n",
    "\n",
    "Since we are using a modified architecture, we wrote the infrastructure ourselves, which may result in some annoyances due to the lack of standardization. Anyway, we load the model as in the code below. We can also generate some text to assess its capabilities. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a bilinear language model, trained on FineWeb-EDU.\n",
    "# This dataset is mostly (high-quality) scientific and educational texts.\n",
    "# This corresponds to GPT2-medium (16 layers, 1024 dimensions).\n",
    "torch.set_grad_enabled(False)\n",
    "model = Transformer.from_pretrained(\"tdooms/fw-medium\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can see the model can generate coherent text.\n",
    "# The model tries to be a smart-ass (FineWeb) but kinda fails.\n",
    "model.generate(\"If I may give you a piece of advice:\", max_length=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then instantiate some useful helper objects: \n",
    "- ``Tracer`` loads SAEs around the MLP of a given layer and contains some helper functions to compute interaction matrices between the two.\n",
    "- ``Visualizer`` shows the (pre-computed) top activations of SAE features to understand their meaning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We set up a Tracer object, which is a utility class to find interesting interactions between two SAEs around an MLP. \n",
    "# Let's inspect a middle layer.\n",
    "tracer = Tracer(model, layer=7, inp=dict(expansion=8), out=dict(expansion=8))\n",
    "\n",
    "# We then create a visualizer for both SAEs.\n",
    "# Implementation-wise, this queries some pre-computed max-activations and shows them in a nice format.\n",
    "inp_vis = Visualizer(model, tracer.inp)\n",
    "out_vis = Visualizer(model, tracer.out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finding interesting output features.\n",
    "With that out of the way, we now get to the interesting bit: analyzing interactions. This can be performed through several means and with multiple levels of rigor; this tutorial will only focus on the simplest: cherry-picking interesting examples. One way we found intriguing examples (that worked well) was by looking at high-level characteristics of all interaction matrices and choosing the outliers. As a metric, we choose the magnitude of the top eigenvectors. Keep in mind that this is simply a heuristic and not extremely principled. Given such features, let's look at their cosine similarities. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running this cell may take awhile depending on your CPU/GPU.\n",
    "\n",
    "# Compute output features whose top eigenvalues are high, likely indicating some interesting structure.\n",
    "eigenvals = tracer.compute(compute_truncated_eigenvalues, project=False, k=2)\n",
    "vals, idxs = eigenvals.topk(10)\n",
    "\n",
    "# Plot the cosine similarity between these features to see if any are related.\n",
    "dirs = tracer.out.w_dec.weight[:, idxs]\n",
    "sims = torch.cosine_similarity(dirs[..., None], dirs[:, None], dim=0)\n",
    "\n",
    "# Visualize them nicely.\n",
    "labels = [f\"{i}\" for i in idxs.cpu()]\n",
    "px.imshow(sims.cpu(), color_continuous_scale=\"RdBu\", color_continuous_midpoint=0, x=labels, y=labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interestingly, the top two features form a (somewhat) linear subspace. We can try to understand their meaning by looking at their top activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's inspect the (somewhat) linear subspace.\n",
    "# This function can visualize arbitrarily many features \n",
    "# It even has a dark mode, which you should disable if you're using a white background.\n",
    "out_vis(3834, 751, dark=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first feature fires on negative words that have recently been negated (not lost, no interference). The second feature does the inverse; it fires on positive words that were negated (not free, little relief). It's not surprising that this forms a linear subspace as the two are opposites, but it is still interesting. What I find fascinating is that we also found this subspace (using the same technique) in a completely different model trained on TinyStories. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Understanding feature interactions.\n",
    "To understand how these features are formed, we can then look at the most salient entries in their interaction matrices. We find that the eigendecomposition of any output feature is quite low-rank. Sadly, while interpreting these eigenvectors in image models is easy (the eigenvectors are simply images), here, the eigenvectors are vectors of all SAE features. Consequently, we have to resort to sparsity.\n",
    "\n",
    "Here, we take the top 50 interactions in the whole interaction matrix and plot their submatrix (the smallest matrix containing all of them). We see that this matrix is quite small, and the features interact in a seemingly systematic manner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the interaction matrix between SAE inputs for a given output SAE feature.\n",
    "q = tracer.q(3834, project=True)\n",
    "\n",
    "# Find the most relevant rows/columns by max interactions.\n",
    "idxs = q.flatten().abs().topk(50).indices\n",
    "i1, i2 = torch.unravel_index(idxs, q.shape)\n",
    "idxs = torch.cat([i1, i2]).unique()\n",
    "\n",
    "# Plot the most important sub-interactions for the output feature.\n",
    "labels = [f\"{i}\" for i in idxs.cpu()]\n",
    "fig = px.imshow(q[idxs.flip(0)][:, idxs].cpu(), color_continuous_midpoint=0, color_continuous_scale=\"RdBu\", x=labels, y=labels, )\n",
    "fig.update_xaxes(side=\"top\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Future directions.\n",
    "While we believe these results are exciting and finally provide a way forward towards understanding MLPs, there is still a lot to explore. We're not at all convinced that the proposed methodology is the best way to understand feature interactions. One possible research avenue is (again) to consider shared structure between features. This can be done using variants of sparse coding or certain tensor decompositions."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
